import jieba
from PIL import Image
from wordcloud import WordCloud, ImageColorGenerator
import numpy as np
from io import BytesIO


# 所有的图片读取后都是bytes类型, Image.open需要进行处理
class WordCloudGenerator:
    """params:
        sentence: 表示词云句子
        source: 接收字节流
        color: 底板颜色
        format: 图片格式
    """
    _colors: tuple = ('white', 'black', 'red', 'blue')

    def __init__(self,
                 sentence: str = 'Hello, word',
                 source: bytes = None,
                 color: str = 'white',
                 format: str = 'JPEG') -> None:
        self._sentence = sentence
        self._source = source
        self._word = None
        self._mask = None
        self._format = format
        self._color = color or 'white'
        # print(self._color)
    def sentence_cut(self) -> None:
        word_list = jieba.cut(sentence=self._sentence)
        self._word = ' '.join(word_list)

    def background(self):
        self._mask = np.array(Image.open(BytesIO(self._source)))

    def cloud(self) -> BytesIO:
        color = ImageColorGenerator(self._mask)

        wordcloud = WordCloud(
            mask=self._mask,
            # TODO: 服务器部署需要改变路径
            font_path='C:/Windows/Fonts/simkai.ttf',
            background_color=self._color,
            # scale=15,
            color_func=color).generate(self._word)
        result = wordcloud.to_image()
        out = BytesIO()
        result.save(out, format=self._format)
        return out

    def create_cloud(self) -> BytesIO:
        self.sentence_cut()
        self.background()
        return self.cloud()


if __name__ == '__main__':
    with open('C:/Users/Fisher/Desktop/a.jpg', 'rb') as f:
        a = f.read()
    generator = WordCloudGenerator("AiSha, Aisha", source=a)
    cloud = generator.create_cloud()
    with open('./a.jpg', 'wb') as f:
        f.write(cloud.getvalue())
